// My_uart.c
// 替换版：基于 UART DMA + IDLE 的帧接收与固件升级处理
// Created by ChatGPT (adapted for your project)

#include "My_uart.h"

#include <stdio.h>
#include <string.h>

#include "main.h"
#include "usart.h"
#include "crc32_table.h"
#include "My_InFlash.h"

// ------------ 协议/缓冲区配置（与原代码一致或可调整） ------------
#define BLOCK_SIZE 4096
#define MAX_FIRMWARE_SIZE (256 * 1024)
#define FRAME_HEADER_LEN 4
#define FRAME_RESERVED_LEN 4
#define FRAME_CRC_LEN 2
#define BLOCK_OVERHEAD 8
#define MAX_FRAME_LEN (FRAME_HEADER_LEN + BLOCK_SIZE + BLOCK_OVERHEAD + FRAME_RESERVED_LEN + FRAME_CRC_LEN)

// 命令码
#define CMD_HANDSHAKE       0x0001
#define CMD_FILE_HEADER     0x0002
#define CMD_DATA_BLOCK      0x0003
#define CMD_FINAL_CHECK     0x0004

// 响应码
#define RESP_SUCCESS        0x00
#define RESP_FAILED         0x01
#define RESP_CRC_ERR        0x02
#define RESP_LEN_ERR        0x03

// 固件信息结构体
typedef struct {
    uint32_t file_size;
    uint32_t total_blocks;
    uint16_t block_size;
    uint32_t global_crc32;
    uint32_t current_block;
    uint32_t recv_size;
    uint32_t crc32_context;
} FirmwareInfo;

// 全局变量（保留你原有名字以最小改动）
FirmwareInfo firmware_info = {0};
uint8_t uart_rx_buf[MAX_FRAME_LEN + 8];    // 用于解析的缓冲区（从 DMA 拷贝过来）
uint16_t uart_rx_len = 0;                  // 接收长度（由 IDLE IRQ 设置）
volatile uint8_t frame_received = 0;       // 帧接收完成标志

static update_callback_uart_t update_Complete_callback = NULL;
static update_callback_uart_t update_Begin_callback = NULL;
static update_callback_uart_t update_Block_callback = NULL;

void set_update_uart_complete_callback(update_callback_uart_t callback) {
    update_Complete_callback = callback;
}
void set_update_uart_Begin_callback(update_callback_uart_t callback) {
    update_Begin_callback = callback;
}
void set_update_uart_Block_callback(update_callback_uart_t callback) {
    update_Block_callback = callback;
}

// ---------------- CRC / helper ----------------
uint16_t crc16_modbus(uint8_t* data, uint16_t len) {
    uint16_t crc = 0xFFFF;
    for (uint16_t i = 0; i < len; i++) {
        crc ^= data[i];
        for (uint8_t j = 0; j < 8; j++) {
            if (crc & 0x0001) {
                crc >>= 1;
                crc ^= 0xA001;
            } else {
                crc >>= 1;
            }
        }
    }
    return crc;
}

uint16_t crc16_xmodem(uint8_t* data, uint16_t len) {
    uint16_t crc = 0;
    for (uint16_t i = 0; i < len; i++) {
        crc ^= (uint16_t)(data[i] << 8);
        for (uint8_t j = 0; j < 8; j++) {
            if (crc & 0x8000) {
                crc = (crc << 1) ^ 0x1021;
            } else {
                crc <<= 1;
            }
        }
    }
    return crc;
}

uint32_t crc32_incremental(uint8_t* data, uint32_t len, uint32_t prev_crc) {
    uint32_t crc = prev_crc ^ 0xFFFFFFFF;
    for (uint32_t i = 0; i < len; i++) {
        crc = (crc >> 8) ^ crc32_table[(crc ^ data[i]) & 0xFF];
    }
    return crc ^ 0xFFFFFFFF;
}

// ---------------- DMA Rx buffer & macros ----------------
// DMA 接收缓冲区：大小至少 >= MAX_FRAME_LEN + 保留
#define DMA_RX_SIZE (MAX_FRAME_LEN + 8)
static uint8_t dma_rx_buf[DMA_RX_SIZE];

// 如果 HAL 没有提供 __HAL_UART_CLEAR_IDLEFLAG，使用 SR/DR 读取清除方法
#ifndef __HAL_UART_CLEAR_IDLEFLAG
#define __HAL_UART_CLEAR_IDLEFLAG(__HANDLE__)  do{ volatile uint32_t tmp; tmp = (__HANDLE__)->Instance->SR; tmp = (__HANDLE__)->Instance->DR; (void)tmp;} while(0)
#endif

// ---------------- UART DMA + IDLE 处理 ----------------
// 启动 DMA 接收并使能 IDLE 中断（在系统初始化或升级模块初始化时调用）
void firmware_upgrade_Init(void) {
    // 启动 DMA 接收（循环缓存/非循环，由 IRQ 停止并重启）
    if (HAL_UART_Receive_DMA(&huart1, dma_rx_buf, DMA_RX_SIZE) != HAL_OK) {
        printf("UART DMA start failed\r\n");
    } else {
        __HAL_UART_ENABLE_IT(&huart1, UART_IT_IDLE); // 使能 IDLE 中断
        printf("UART DMA RX started, IDLE IRQ enabled\r\n");
    }
}

// IDLE 检测处理函数（在 USARTx_IRQHandler 中调用）
static void uart_idle_process(void) {
    // 获取 DMA 剩余计数（hdmarx 必须有效）
    if (huart1.hdmarx == NULL) {
        // 无 DMA handle，直接返回（不应发生）
        return;
    }

    // 先清 IDLE 标志（通过宏或读 SR/DR）
    __HAL_UART_CLEAR_IDLEFLAG(&huart1);

    // 得到 DMA 剩余计数（注意 HAL 宏）
    uint32_t dma_remain = __HAL_DMA_GET_COUNTER(huart1.hdmarx);
    uint16_t received = (uint16_t)(DMA_RX_SIZE - dma_remain);

    // 停止 DMA 以便安全读取 buffer
    HAL_UART_DMAStop(&huart1);


    if (received == 0 || received > (MAX_FRAME_LEN + 8)) {
        // 长度异常：忽略并重启 DMA
        uart_rx_len = 0;
        frame_received = 0;
        // 重新启动 DMA
        HAL_UART_Receive_DMA(&huart1, dma_rx_buf, DMA_RX_SIZE);
        return;
    }

    // 拷贝到解析缓冲区
    memcpy(uart_rx_buf, dma_rx_buf, received);
    uart_rx_len = received;
    frame_received = 1;

    // 立即重启 DMA 接收准备下一帧
    HAL_UART_Receive_DMA(&huart1, dma_rx_buf, DMA_RX_SIZE);
}

// 注意：如果你的工程已经实现了USART1_IRQHandler，请把下面的实现合并到你的 IRQ 文件中。
// 直接替换可能导致重复定义：如果冲突，把内部 IDLE-check 部分搬到你的 IRQ 中并调用 uart_idle_process()。
void USART1_IRQHandler(void)
{
    // 检查 IDLE
    if (__HAL_UART_GET_FLAG(&huart1, UART_FLAG_IDLE)) {
        uart_idle_process();
    }
    // 调用 HAL 的 IRQ Handler 处理其它事件（TX complete 等）
    HAL_UART_IRQHandler(&huart1);
}

// ---------------------------- 帧解析与处理（保留并略微调整） ----------------------------
// 解析帧数据（返回命令码和数据区）
uint8_t parse_frame(uint16_t* cmd, uint8_t** data, uint16_t* data_len) {
    if (uart_rx_len < (FRAME_HEADER_LEN + FRAME_RESERVED_LEN + FRAME_CRC_LEN)) {
        return RESP_LEN_ERR;
    }
    // CRC（帧头+数据区+保留字段）
    uint16_t frame_crc = (uint16_t)(uart_rx_buf[uart_rx_len - 1] << 8 | uart_rx_buf[uart_rx_len - 2]);
    uint16_t calc_crc = crc16_modbus(uart_rx_buf, uart_rx_len - FRAME_CRC_LEN);
    if (frame_crc != calc_crc) {
        return RESP_CRC_ERR;
    }
    // 命令码（小端：低字节在 uart_rx_buf[0]）
    *cmd = (uint16_t)(uart_rx_buf[1] << 8 | uart_rx_buf[0]);
    // 数据长度（小端）
    *data_len = (uint16_t)(uart_rx_buf[3] << 8 | uart_rx_buf[2]);
    *data = &uart_rx_buf[FRAME_HEADER_LEN];
    // 检查长度匹配
    if (*data_len != (uart_rx_len - FRAME_HEADER_LEN - FRAME_RESERVED_LEN - FRAME_CRC_LEN)) {
        return RESP_LEN_ERR;
    }
    return RESP_SUCCESS;
}

// 发送响应帧（不在中断中调用）
void send_response(uint16_t cmd, uint8_t resp_code) {
    uint8_t resp_data = resp_code;
    uint8_t frame[FRAME_HEADER_LEN + 1 + FRAME_RESERVED_LEN + FRAME_CRC_LEN];

    frame[0] = (uint8_t)(cmd & 0xFF);
    frame[1] = (uint8_t)(cmd >> 8);
    frame[2] = (uint8_t)(1 & 0xFF);
    frame[3] = (uint8_t)(1 >> 8);

    frame[4] = resp_data;

    frame[5] = 0x00;
    frame[6] = 0x00;
    frame[7] = 0x00;
    frame[8] = 0x00;

    uint16_t crc = crc16_modbus(frame, 9);
    frame[9] = (uint8_t)(crc & 0xFF);
    frame[10] = (uint8_t)(crc >> 8);

    HAL_UART_Transmit(&huart1, frame, 11, 200); // 适当延长超时时间
}

// ---------------- 升级命令处理（几乎照搬你原来的实现） ----------------
void handle_handshake(uint8_t* data, uint16_t data_len) {
    if (data_len == 1 && data[0] == 0x01) {
        firmware_info.current_block = 0;
        firmware_info.recv_size = 0;
        firmware_info.crc32_context = 0;
        send_response(CMD_HANDSHAKE, RESP_SUCCESS);
    } else {
        send_response(CMD_HANDSHAKE, RESP_FAILED);
    }
}

void handle_file_header(uint8_t* data, uint16_t data_len) {
    if (data_len != 4 + 4 + 2 + 4) {
        send_response(CMD_FILE_HEADER, RESP_LEN_ERR);
        return;
    }
    firmware_info.file_size = *(uint32_t*)&data[0];
    firmware_info.total_blocks = *(uint32_t*)&data[4];
    firmware_info.block_size = *(uint16_t*)&data[8];
    firmware_info.global_crc32 = *(uint32_t*)&data[10];

    if (firmware_info.file_size > MAX_FIRMWARE_SIZE ||
        firmware_info.block_size != BLOCK_SIZE ||
        firmware_info.total_blocks == 0) {
        send_response(CMD_FILE_HEADER, RESP_FAILED);
        return;
    }

    // 注意：这里保留原有直接擦除逻辑，但在 DMA 模式下擦除期间请确保不会禁中断或影响 UART。
    // 更稳妥的做法是把擦除作为单独的命令（CMD_ERASE），或把擦除函数放到 RAM 执行。
    if (InternalFlash_Erase(APP2, 256 * 1024, 2)) {
        printf("擦除失败\r\n");
        send_response(CMD_FILE_HEADER, RESP_FAILED);
        return;
    }

    if (update_Begin_callback) {
        update_Begin_callback(1);
    }
    send_response(CMD_FILE_HEADER, RESP_SUCCESS);
}

void handle_data_block(uint8_t* data, uint16_t data_len) {
    if (data_len < 4 + 2 + 2) {
        send_response(CMD_DATA_BLOCK, RESP_LEN_ERR);
        return;
    }
    uint32_t block_idx = *(uint32_t*)&data[0];
    uint16_t block_len = *(uint16_t*)&data[4];
    uint8_t* block_data = &data[6];
    uint16_t block_crc16 = *(uint16_t*)&data[6 + block_len];

    if (block_idx != firmware_info.current_block) {
        send_response(CMD_DATA_BLOCK, RESP_FAILED);
        return;
    }
    if (block_len > firmware_info.block_size ||
        (6 + block_len + 2) != data_len) {
        send_response(CMD_DATA_BLOCK, RESP_LEN_ERR);
        return;
    }
    if (crc16_xmodem(block_data, block_len) != block_crc16) {
        send_response(CMD_DATA_BLOCK, RESP_CRC_ERR);
        return;
    }

    if (update_Block_callback) {
        update_Block_callback(1);
    }

    // 实际写入（请确保 InFlash_Write 实现不会长时间禁用中断）
    uint8_t err = InFlash_Write(APP2 + (block_idx * (uint32_t)firmware_info.block_size), block_data, block_len);
    if (err) {
        printf("InFlash_Write err : %d\r\n", err);
        send_response(CMD_DATA_BLOCK, RESP_FAILED);
        return;
    }

    firmware_info.recv_size += block_len;
    firmware_info.crc32_context = crc32_incremental(block_data, block_len, firmware_info.crc32_context);

    firmware_info.current_block++;
    send_response(CMD_DATA_BLOCK, RESP_SUCCESS);
}

void handle_final_check(uint8_t* data, uint16_t data_len) {
    if (firmware_info.recv_size != firmware_info.file_size ||
        firmware_info.current_block != firmware_info.total_blocks) {
        send_response(CMD_FINAL_CHECK, RESP_FAILED);
        return;
    }
    if (firmware_info.crc32_context != firmware_info.global_crc32) {
        send_response(CMD_FINAL_CHECK, RESP_CRC_ERR);
        if (update_Complete_callback) update_Complete_callback(1);
        return;
    }
    if (update_Complete_callback) update_Complete_callback(0);
    send_response(CMD_FINAL_CHECK, RESP_SUCCESS);
}

// ---------------- 主循环处理 ----------------
void firmware_upgrade_Loop(void) {
    if (frame_received) {
        // 先把标志清掉避免重入（DMA 已被重新启动）
        frame_received = 0;

        uint16_t cmd;
        uint8_t* data;
        uint16_t data_len;
        uint8_t parse_result = parse_frame(&cmd, &data, &data_len);

        if (parse_result != RESP_SUCCESS) {
            // 若 parse_frame 做不到提供 cmd（如长度错误），我们不能用 cmd 来回送
            // 采用通用失败回复（或直接忽略）
            // 这里如果 parse_frame 返回 RESP_LEN_ERR，我们可以尝试不回 cmd（因为 cmd 未解析）
            if (parse_result == RESP_LEN_ERR) {
                // 直接打印，方便调试
                printf("parse_frame: LEN_ERR uart_rx_len=%u\r\n", uart_rx_len);
            } else if (parse_result == RESP_CRC_ERR) {
                printf("parse_frame: CRC_ERR uart_rx_len=%u\r\n", uart_rx_len);
            }
            // 若能取 cmd，则回复
            // safe extraction: if uart_rx_len >=2 then extract cmd bytes
            if (uart_rx_len >= 2) {
                uint16_t fake_cmd = (uint16_t)(uart_rx_buf[1] << 8 | uart_rx_buf[0]);
                send_response(fake_cmd, parse_result);
            }
        } else {
            switch (cmd) {
                case CMD_HANDSHAKE:
                    handle_handshake(data, data_len);
                    break;
                case CMD_FILE_HEADER:
                    handle_file_header(data, data_len);
                    break;
                case CMD_DATA_BLOCK:
                    handle_data_block(data, data_len);
                    break;
                case CMD_FINAL_CHECK:
                    handle_final_check(data, data_len);
                    break;
                default:
                    send_response(cmd, RESP_FAILED);
                    break;
            }
        }

        // 解析完毕，复位接收长度（DMA 已在 IRQ 中自动重启）
        uart_rx_len = 0;
    }
}
